package db

import (
	"fmt"
	lru "github.com/hashicorp/golang-lru/v2"
	"gorm.io/driver/mysql"
	"gorm.io/driver/postgres"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
	"strings"
	"sync"
	"time"
	"zatta-gateway/pkg/db/model"
	"zatta-gateway/pkg/grpc"
	"zatta-gateway/pkg/utils"
)

type ClientManager struct {
	// database managers
	modelDBDNS   []string
	rpcServerDNS []string
	modelDBs     []*gorm.DB
	rpcServerDBs []*gorm.DB

	// 负载均衡策略
	dbIndexMu        sync.RWMutex
	modelDBIndex     int
	rpcServerDBIndex int

	// rpc client manager
	rpcClients map[string]*grpc.Backend // key: serviceAddr
	clientsMu  sync.RWMutex

	// cache for query results
	queryCache *lru.Cache[string, interface{}]
	cacheMu    sync.RWMutex
}

func NewClientManager(modelDBDNS, rpcServerDNS []string) *ClientManager {
	queryCache, _ := lru.New[string, interface{}](1000)
	return &ClientManager{
		modelDBDNS:   modelDBDNS,
		rpcServerDNS: rpcServerDNS,
		rpcClients:   make(map[string]*grpc.Backend),
		queryCache:   queryCache,
		modelDBs:     make([]*gorm.DB, 0, len(modelDBDNS)),
		rpcServerDBs: make([]*gorm.DB, 0, len(rpcServerDNS)),
	}
}

func (cm *ClientManager) Connect() error {
	// 连接所有 model DBs
	for _, dns := range cm.modelDBDNS {
		modelDB, err := cm.connectDB(dns)
		if err != nil {
			return fmt.Errorf("failed to connect to model database (%s): %w", dns, err)
		}
		cm.modelDBs = append(cm.modelDBs, modelDB)
	}

	// 连接所有 rpc server DBs
	for _, dns := range cm.rpcServerDNS {
		rpcServerDB, err := cm.connectDB(dns)
		if err != nil {
			return fmt.Errorf("failed to connect to rpc server database (%s): %w", dns, err)
		}
		cm.rpcServerDBs = append(cm.rpcServerDBs, rpcServerDB)
	}

	return nil
}
func (cm *ClientManager) connectDB(dns string) (*gorm.DB, error) {
	var dialector gorm.Dialector
	switch {
	case strings.HasPrefix(dns, "mysql://"):
		dialector = mysql.Open(strings.TrimPrefix(dns, "mysql://"))
	case strings.HasPrefix(dns, "postgres://"):
		dialector = postgres.Open(dns)
	case strings.HasPrefix(dns, "sqlite://"):
		dialector = sqlite.Open(strings.TrimPrefix(dns, "sqlite://"))
	default:
		// 默认当作 MySQL 处理（移除前缀）
		if strings.Contains(dns, "@tcp(") {
			dialector = mysql.Open(dns)
		} else {
			return nil, fmt.Errorf("unsupported database type or invalid DNS format")
		}
	}

	db, err := gorm.Open(dialector, &gorm.Config{})
	if err != nil {
		return nil, fmt.Errorf("failed to connect to database: %w", err)
	}

	sqlDB, err := db.DB()
	if err != nil {
		return nil, fmt.Errorf("failed to get database instance: %w", err)
	}

	// 设置连接池参数
	sqlDB.SetMaxIdleConns(20)
	sqlDB.SetMaxOpenConns(200)
	sqlDB.SetConnMaxLifetime(30 * time.Minute)
	sqlDB.SetConnMaxIdleTime(10 * time.Minute)

	return db, nil
}

// roundRobinModelDB 使用轮询策略选择model数据库
func (cm *ClientManager) roundRobinModelDB() *gorm.DB {
	cm.dbIndexMu.Lock()
	defer cm.dbIndexMu.Unlock()

	if cm.GetModelDBCount() == 0 {
		return nil
	}
	if cm.GetModelDBCount() == 1 {
		return cm.modelDBs[0]
	}
	db := cm.modelDBs[cm.modelDBIndex]
	cm.modelDBIndex = utils.RoundRobinMode(cm.modelDBIndex, len(cm.modelDBs))
	return db
}

// roundRobinRpcServerDB 使用轮询策略选择rpc server数据库
func (cm *ClientManager) roundRobinRpcServerDB() *gorm.DB {
	cm.dbIndexMu.Lock()
	defer cm.dbIndexMu.Unlock()

	if cm.GetRpcServerDBCount() == 0 {
		return nil
	}
	if cm.GetRpcServerDBCount() == 1 {
		return cm.rpcServerDBs[0]
	}
	db := cm.rpcServerDBs[cm.rpcServerDBIndex]
	cm.rpcServerDBIndex = utils.RoundRobinMode(cm.rpcServerDBIndex, len(cm.rpcServerDBs))
	return db
}

// GetModelDB 获取模型数据库连接（使用轮询负载均衡）
func (cm *ClientManager) GetModelDB() *gorm.DB {
	return cm.roundRobinModelDB()
}

// GetRpcServerDB 获取RPC服务器数据库连接（使用轮询负载均衡）
func (cm *ClientManager) GetRpcServerDB() *gorm.DB {
	return cm.roundRobinRpcServerDB()
}

// GetRpcServersByModelName 从rpcServerDB获取指定模型的RPC服务器列表
func (cm *ClientManager) GetRpcServersByModelName(modelName string) ([]*model.RpcServer, error) {
	cacheKey := "rpc_servers_" + modelName

	// 尝试从缓存获取
	if cached, ok := cm.getCache(cacheKey); ok {
		if servers, ok := cached.([]*model.RpcServer); ok {
			return servers, nil
		}
	}

	// 缓存未命中，查询数据库
	// 使用轮询策略选择数据库
	db := cm.GetRpcServerDB()
	if db == nil {
		return nil, fmt.Errorf("no available rpc server database")
	}

	rpcServers, err := model.ListRpcServerByModelName(db, modelName)
	if err != nil {
		return nil, err
	}

	// 将结果存入缓存
	cm.setCache(cacheKey, rpcServers)

	return rpcServers, nil
}

// GetSpecificModelDB 获取特定索引的model数据库连接
func (cm *ClientManager) GetSpecificModelDB(index int) (*gorm.DB, error) {
	if index < 0 || index >= len(cm.modelDBs) {
		return nil, fmt.Errorf("invalid model db index: %d", index)
	}
	return cm.modelDBs[index], nil
}

// GetSpecificRpcServerDB 获取特定索引的rpc server数据库连接
func (cm *ClientManager) GetSpecificRpcServerDB(index int) (*gorm.DB, error) {
	if index < 0 || index >= len(cm.rpcServerDBs) {
		return nil, fmt.Errorf("invalid rpc server db index: %d", index)
	}
	return cm.rpcServerDBs[index], nil
}

// GetModelDBCount 获取model数据库数量
func (cm *ClientManager) GetModelDBCount() int {
	return len(cm.modelDBs)
}

// GetRpcServerDBCount 获取rpc server数据库数量
func (cm *ClientManager) GetRpcServerDBCount() int {
	return len(cm.rpcServerDBs)
}

// GetRpcBackend 根据服务地址获取或创建RPC后端客户端
func (cm *ClientManager) GetRpcBackend(serviceAddr string) grpc.Backend {
	var client *grpc.Backend
	var exists bool
	cm.clientsMu.RLock()
	client, exists = cm.rpcClients[serviceAddr]
	cm.clientsMu.RUnlock()

	if exists {
		return *client
	}

	// 创建新的RPC客户端
	cm.clientsMu.Lock()
	defer cm.clientsMu.Unlock()

	// 双重检查
	if client, exists = cm.rpcClients[serviceAddr]; exists {
		return *client
	}

	// 创建新的gRPC后端
	backend := grpc.NewClient(serviceAddr, false, nil, false)
	cm.rpcClients[serviceAddr] = &backend

	return backend
}

// Cache操作的封装方法
func (cm *ClientManager) getCache(key string) (interface{}, bool) {
	cm.cacheMu.RLock()
	defer cm.cacheMu.RUnlock()
	return cm.queryCache.Get(key)
}

func (cm *ClientManager) setCache(key string, value interface{}) {
	cm.cacheMu.Lock()
	defer cm.cacheMu.Unlock()
	cm.queryCache.Add(key, value)
}

func (cm *ClientManager) removeCache(key string) {
	cm.cacheMu.Lock()
	defer cm.cacheMu.Unlock()
	cm.queryCache.Remove(key)
}
